package com.supreme.nettyapp.generalserver.mqttserver.handler;

import com.supreme.nettyapp.generalserver.mqttserver.model.SessionStore;
import com.supreme.nettyapp.generalserver.utils.ChannelStoreUtils;
import io.netty.buffer.Unpooled;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.SimpleChannelInboundHandler;
import io.netty.handler.codec.DecoderResult;
import io.netty.handler.codec.mqtt.*;
import io.netty.util.AttributeKey;

import java.nio.charset.StandardCharsets;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Random;
import java.util.function.Function;
import java.util.stream.Collectors;

@ChannelHandler.Sharable
public class MqttServerHandler extends SimpleChannelInboundHandler<MqttMessage> {

    private static final AttributeKey<String> CLIENT_ID = AttributeKey.valueOf("clientId");

    @Override
    protected void channelRead0(ChannelHandlerContext ctx, MqttMessage msg) throws Exception {
        MqttFixedHeader mqttFixedHeader = msg.fixedHeader();
        MqttMessageType mqttMessageType = mqttFixedHeader.messageType();
        switch (mqttMessageType) {
            case CONNECT:
                this.handleConnect(ctx, (MqttConnectMessage) msg);
                break;

            case PUBLISH:
                this.handlePublish(ctx, (MqttPublishMessage) msg);
                break;

            case SUBSCRIBE:
                this.handleSubscribe(ctx, (MqttSubscribeMessage) msg);
                break;

            case UNSUBSCRIBE:
                this.handleUnsubscribe(ctx, (MqttUnsubscribeMessage) msg);
                break;

            case DISCONNECT:
                this.handleDisconnect(ctx);
                break;

            case PINGREQ:
                this.handlePingReq(ctx);
                break;

            default:
                break;
        }
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        cause.printStackTrace();
        ctx.close();
    }

    /**
     * 这份代码实现了处理MQTT Connect消息的具体代码，当客户端发送Connect消息时，服务端将向客户端发送ConnAck消息，以确认连接请求是否被接受。
     * @param ctx
     * @param msg
     */
    private void handleConnect(ChannelHandlerContext ctx, MqttConnectMessage msg) {
        DecoderResult decoderResult = msg.decoderResult();
        boolean failure = decoderResult.isFailure();
        if (failure) {
            Throwable cause = decoderResult.cause();
            if (cause instanceof MqttUnacceptableProtocolVersionException) {
                // 不支持的协议版本
                MqttConnAckMessage connAckMessage = (MqttConnAckMessage) MqttMessageFactory.newMessage(
                        new MqttFixedHeader(
                                MqttMessageType.CONNACK,
                                false,
                                MqttQoS.AT_MOST_ONCE,
                                false,
                                0),
                        new MqttConnAckVariableHeader(MqttConnectReturnCode.CONNECTION_REFUSED_UNACCEPTABLE_PROTOCOL_VERSION, false),
                        null);
                ctx.writeAndFlush(connAckMessage);
            } else if (cause instanceof MqttIdentifierRejectedException) {
                // 不合格的clientId
                MqttConnAckMessage connAckMessage = (MqttConnAckMessage) MqttMessageFactory.newMessage(
                        new MqttFixedHeader(
                                MqttMessageType.CONNACK,
                                false,
                                MqttQoS.AT_MOST_ONCE,
                                false,
                                0),
                        new MqttConnAckVariableHeader(MqttConnectReturnCode.CONNECTION_REFUSED_IDENTIFIER_REJECTED, false),
                        null);
                ctx.writeAndFlush(connAckMessage);
            }

            ctx.close();
            return;
        }

        MqttConnectPayload payload = msg.payload();
        String clientIdentifier = payload.clientIdentifier();
        String userName = payload.userName();
        String password = Optional.ofNullable(payload.passwordInBytes()).map(String::new).orElse(null);
        System.out.printf("CONNECT clientId = %s username = %s password = %s\n", clientIdentifier, userName, password);

        if (Objects.isNull(clientIdentifier)) {
            MqttConnAckMessage connAckMessage = (MqttConnAckMessage) MqttMessageFactory.newMessage(
                    new MqttFixedHeader(
                            MqttMessageType.CONNACK,
                            false,
                            MqttQoS.AT_MOST_ONCE,
                            false,
                            0),
                    new MqttConnAckVariableHeader(MqttConnectReturnCode.CONNECTION_REFUSED_IDENTIFIER_REJECTED, false),
                    null);
            ctx.writeAndFlush(connAckMessage);
            ctx.close();
            return;
        }

        MqttConnAckMessage connAckMessage = (MqttConnAckMessage) MqttMessageFactory.newMessage(
                new MqttFixedHeader(
                        MqttMessageType.CONNACK,
                        false,
                        MqttQoS.AT_MOST_ONCE,
                        false,
                        0),
                new MqttConnAckVariableHeader(MqttConnectReturnCode.CONNECTION_ACCEPTED, false),
                null);

        Channel channel = ctx.channel();
        channel.attr(CLIENT_ID).set(clientIdentifier);

        SessionStore sessionStore = new SessionStore()
                .setClientId(clientIdentifier)
                .setUsername(userName)
                .setPassword(password)
                .setChannel(channel);
        ChannelStoreUtils.putSessionStore(clientIdentifier, sessionStore);

        ctx.writeAndFlush(connAckMessage);
    }

    /**
     * 这份代码实现了处理MQTT Subscribe消息的具体代码，当客户端发送Subscribe消息时，服务端将该客户端订阅的主题加入到主题的客户端列表中。并且回复一条SubAck
     * @param ctx
     * @param msg
     */
    private void handleSubscribe(ChannelHandlerContext ctx, MqttSubscribeMessage msg) {
        List<MqttTopicSubscription> topicSubscriptions = msg.payload().topicSubscriptions();
        Channel channel = ctx.channel();
        String clientId = channel.attr(CLIENT_ID).get();
        List<String> topicList = topicSubscriptions.stream()
                .map(MqttTopicSubscription::topicName)
                .collect(Collectors.toList());
        System.out.printf("SUBSCRIBE clientId = %s topicList = %s\n", clientId, topicList);

        for (MqttTopicSubscription topicSubscription : topicSubscriptions) {
            String topic = topicSubscription.topicName();
            ChannelStoreUtils.addSessionTopic(clientId, topic);
        }

        MqttFixedHeader mqttFixedHeader = new MqttFixedHeader(
                MqttMessageType.SUBACK,
                false,
                /*mqttQoS*/MqttQoS.AT_MOST_ONCE,
                false,
                0);
        MqttMessageIdVariableHeader mqttMessageIdVariableHeader = MqttMessageIdVariableHeader.from(msg.variableHeader().messageId());

        MqttSubAckPayload mqttSubAckPayload = new MqttSubAckPayload(
                topicSubscriptions.stream()
                        .map(topicSubscription -> topicSubscription.qualityOfService().value())
                        .collect(Collectors.toList()));
        MqttSubAckMessage mqttSubAckMessage = new MqttSubAckMessage(mqttFixedHeader, mqttMessageIdVariableHeader, mqttSubAckPayload);

        ctx.writeAndFlush(mqttSubAckMessage);
    }

    /**
     * 这份代码实现了处理MQTT Unsubscribe消息的具体代码，当客户端发送Unsubscribe消息时，服务端将该客户端取消订阅的主题从主题的客户端列表中移除。并且回复一条UnsubAck消息。
     * @param ctx
     * @param msg
     */
    private void handleUnsubscribe(ChannelHandlerContext ctx, MqttUnsubscribeMessage msg) {
        List<String> topicFilters = msg.payload().topics();
        Channel channel = ctx.channel();
        String clientId = channel.attr(CLIENT_ID).get();
        for (String topicFilter : topicFilters) {
            ChannelStoreUtils.removeSessionTopic(clientId, topicFilter);
        }
        MqttFixedHeader mqttFixedHeader = new MqttFixedHeader(
                MqttMessageType.UNSUBACK,
                false,
                MqttQoS.AT_LEAST_ONCE,
                false,
                0);
        MqttMessageIdVariableHeader mqttMessageIdVariableHeader = MqttMessageIdVariableHeader.from(msg.variableHeader().messageId());
        MqttUnsubAckMessage mqttUnsubAckMessage = new MqttUnsubAckMessage(mqttFixedHeader, mqttMessageIdVariableHeader);
        ctx.writeAndFlush(mqttUnsubAckMessage);
    }

    /**
     * 这份代码实现了处理MQTT DISCONNECT消息的具体代码，当客户端发送DISCONNECT消息时，服务端将该客户端所订阅的所有主题从主题的客户端列表中移除，并且关闭该客户端的通道。
     * @param ctx
     */
    private void handleDisconnect(ChannelHandlerContext ctx) {
        String clientId = ctx.channel().attr(CLIENT_ID).get();
        if (clientId != null) {
            ChannelStoreUtils.removeSessionStore(clientId);
            ctx.channel().attr(CLIENT_ID).set(null);
        }
        ctx.close();
    }

    /**
     * 这份代码实现了处理MQTT Publish消息的具体代码，当客户端发送Publish消息时，服务端将该消息转发给所有订阅了该主题的客户端。
     * @param ctx
     * @param msg
     */
    private void handlePublish(ChannelHandlerContext ctx, MqttPublishMessage msg) {
        Channel channel = ctx.channel();
        String clientId = channel.attr(CLIENT_ID).get();
        if (clientId != null) {
            String topic = msg.variableHeader().topicName();
            MqttQoS mqttQoS = msg.fixedHeader().qosLevel();

            if (ChannelStoreUtils.containSessionTopic(clientId, topic)) {
                byte[] payload = new byte[msg.payload().readableBytes()];
                msg.payload().getBytes(msg.payload().readerIndex(), payload);
                String payloadContent = new String(payload, StandardCharsets.UTF_8);
                System.out.printf("PUBLISH clientId = %s content = %s\n", clientId, payloadContent);

                int packetId = msg.variableHeader().packetId();
                MqttPubAckMessage mqttPubAckMessage = (MqttPubAckMessage) MqttMessageFactory.newMessage(
                        new MqttFixedHeader(
                                MqttMessageType.PUBACK,
                                false,
                                mqttQoS,
                                false,
                                0),
                        MqttMessageIdVariableHeader.from(new Random().nextInt(65536)),
                        null);
                channel.writeAndFlush(mqttPubAckMessage);

                MqttPublishMessage mqttPublishMessage = (MqttPublishMessage) MqttMessageFactory.newMessage(
                        new MqttFixedHeader(
                                MqttMessageType.PUBLISH,
                                false,
                                mqttQoS,
                                false,
                                0),
                        new MqttPublishVariableHeader(topic, packetId),
                        Unpooled.wrappedBuffer(payload)
                );
                channel.writeAndFlush(mqttPublishMessage);
            }
        }
    }

    /**
     * 这份代码实现了处理MQTT PINGREQ消息的具体代码，当客户端发送PINGREQ消息时，服务端会回复一条PINGRESP消息。这主要用于维护客户端和服务端的连接，保证客户端与服务端的连接是有效的。
     * @param ctx
     */
    private void handlePingReq(ChannelHandlerContext ctx) {
        MqttFixedHeader mqttFixedHeader = new MqttFixedHeader(
                MqttMessageType.PINGRESP,
                false,
                MqttQoS.AT_MOST_ONCE,
                false,
                0);
        MqttMessage mqttMessage = new MqttMessage(mqttFixedHeader);
        ctx.writeAndFlush(mqttMessage);
    }
}
